from __future__ import annotations

from typing import TYPE_CHECKING

import gymnasium as gym
import numpy as np
from ray.rllib.models.modelv2 import restore_original_dimensions
from ray.rllib.models.preprocessors import get_preprocessor

if TYPE_CHECKING:
    from driving_gym.environment import DrivingGym


class FlattenSpaceWrapper(gym.Wrapper):
    """
    Convert Dict observation space to Box observation space. Suitable for d3rlpy.
    Also provides functionality to restore the original observation structure.
    """

    def __init__(self, env: DrivingGym):
        super().__init__(env)

        # Store the original observation space
        self.original_space = env.observation_space

        # Flatten the observation space
        sample = env.observation_space.sample()
        self.preprossor = get_preprocessor(env.observation_space)(env.observation_space)
        processed = self.preprossor.transform(sample)
        self.observation_space = gym.spaces.Box(
            low=-np.inf, high=np.inf, shape=processed.shape, dtype=processed.dtype
        )

    def reset(self, *, seed=None, options=None):
        obs, info = self.env.reset(seed=seed, options=options)
        return self.preprossor.transform(obs), info

    def step(self, action):
        obs, reward, te, tr, info = self.env.step(action)
        return self.preprossor.transform(obs), reward, te, tr, info

    @staticmethod
    def restore_observation(obs, original_space, tensorlib="numpy"):
        return restore_original_dimensions(obs, original_space, tensorlib)
